import torch
import numpy as np

def default_collate(dataset_dicts):
    batch_size = len(dataset_dicts)
    image_list = [dataset_dict['image'] for dataset_dict in dataset_dicts]
    
    category_ids_list = [
        torch.FloatTensor([object['category_id'] for object in dataset_dict['annotations']])
        for dataset_dict in dataset_dicts]
    
    bboxes_list = [
        torch.FloatTensor([object['bbox'] for object in dataset_dict['annotations']]).view(-1, 4)
        for dataset_dict in dataset_dicts]

    image_list = torch.from_numpy(np.stack(image_list, axis=0))

    """
    max_num_objects = max(len(category_ids) for category_ids in category_ids_list)
    
    if max_num_objects > 0:
        category_ids_list_padded = torch.ones((batch_size, max_num_objects, 4)) * -1
        bboxes_list_padded = torch.ones((batch_size, max_num_objects)) * -1

        for i, category_ids, bboxes in enumerate(zip(category_ids_list, bboxes_list)):
            category_ids_list_padded[i, :len(category_ids), :] = category_ids
            bboxes_list_padded[i, :len(category_ids)] = bboxes
    else:
        category_ids_list_padded = torch.ones((batch_size, 0, 4))
        bboxes_list_padded = torch.ones((batch_size, 0))
    """

    return {'image': image_list, 'annotations': {'category_ids': category_ids_list, 'bboxes': bboxes_list}}
    
